package cas.ihep.fs.alluxio;

import alluxio.Configuration;
import alluxio.PropertyKey;
import alluxio.client.file.options.InStreamOptions;
import alluxio.client.file.options.OutStreamOptions;
import alluxio.proto.dataserver.Protocol;
import alluxio.util.proto.ProtoMessage;
import alluxio.wire.WorkerNetAddress;
import cas.ihep.fs.UnifiedBlock;
import cas.ihep.fs.UnifiedFile;
import cas.ihep.util.AlluxioRPCContext;
import cas.ihep.util.Finalizer;
import com.google.common.base.Preconditions;

import java.io.Closeable;
import java.io.File;
import java.io.IOException;
import java.lang.management.ManagementFactory;
import java.lang.management.RuntimeMXBean;
import java.util.concurrent.atomic.AtomicInteger;

abstract class AlluxioBlock implements UnifiedBlock,Closeable{
    Finalizer closer;
    String blockPath;
    long thisId,thisSize;
    boolean isClosed;
    AlluxioRPCContext rpcContext;

    AlluxioBlock(long id_,long size_){
        closer=Finalizer.create();
        blockPath=null;
        thisId=id_;
        thisSize=size_;
        isClosed=false;
    }

    AlluxioBlock manageRPC(){
        closer.registerLast(rpcContext);
        return this;
    }

    public long id(){
        return thisId;
    }

    public long size(){
        return thisSize;
    }

    public String absolutePath() {
        return blockPath;
    }

    static class AlluxioReadableBlock extends AlluxioBlock{

        private static final long READ_TIMEOUT_MS =
                Configuration.getMs(PropertyKey.USER_NETWORK_NETTY_TIMEOUT_MS);

        ///private ClosableResource mResource;

        AlluxioReadableBlock(AlluxioFile.AlluxioFileIn parent,
                             WorkerNetAddress address,long id_, long size_) throws IOException{
            super(id_, size_);
            try {
                AlluxioFileSystem fileSystem=parent.alluxioFileSystem;
                rpcContext=closer.register(new AlluxioRPCContext(fileSystem.fileSystemContext,address));
                ///mResource=super.closer.register(new ClosableResource(fileSystem.fileSystemContext,address));
                Protocol.LocalBlockOpenRequest request =
                        Protocol.LocalBlockOpenRequest.newBuilder().setBlockId(id_)
                                .setPromote(parent.mInOptions.getAlluxioStorageType().isPromote()).build();
                ///ProtoMessage message = NettyRPC
                ///        .call(NettyRPCContext.defaults().setChannel(mResource.channel).setTimeout(READ_TIMEOUT_MS),
                ///                new ProtoMessage(request));
                ProtoMessage message= rpcContext.call(new ProtoMessage(request),READ_TIMEOUT_MS);
                Preconditions.checkState(message.isLocalBlockOpenResponse());
                super.blockPath = message.asLocalBlockOpenResponse().getPath();
            } catch(IOException e) {
                try {
                    super.closer.close();
                }finally {
                    super.closer=null;
                }
                throw e;
            }
        }

        AlluxioReadableBlock(AlluxioRPCContext rpc, long id_, long size_, InStreamOptions mInOptions)throws IOException{
            super(id_,size_);
            rpcContext=rpc;
            Protocol.LocalBlockOpenRequest request =
                    Protocol.LocalBlockOpenRequest.newBuilder().setBlockId(id_)
                            .setPromote(mInOptions.getAlluxioStorageType().isPromote()).build();
            ProtoMessage message= rpc.call(new ProtoMessage(request),READ_TIMEOUT_MS);
            Preconditions.checkState(message.isLocalBlockOpenResponse());
            blockPath = message.asLocalBlockOpenResponse().getPath();
        }

        @Override
        public void cancel() throws IOException {
            close();
        }

        @Override
        public int permission() {
            return UnifiedFile.PERMISSION_READ;
        }

        @Override
        public void close() throws IOException {
            if(!isClosed) {
                isClosed=true;
                try{
                    Protocol.LocalBlockCloseRequest request =
                            Protocol.LocalBlockCloseRequest.newBuilder().setBlockId(thisId).build();
                    ///NettyRPC.call(NettyRPCContext.defaults().setChannel(mResource.channel).setTimeout(READ_TIMEOUT_MS),
                    ///        new ProtoMessage(request));
                    rpcContext.call(new ProtoMessage(request));
                }finally {
                    closer.close();
                }
            }
        }
    }

    private static final long FILE_BUFFER_BYTES =
            Configuration.getBytes(PropertyKey.USER_FILE_BUFFER_BYTES);
    private static final long WRITE_TIMEOUT_MS =
            Configuration.getMs(PropertyKey.USER_NETWORK_NETTY_TIMEOUT_MS);

    static class AlluxioWritableBlock2 extends AlluxioBlock{
        ///private NettyRPCContext mNettyRPCContext;
        private ProtoMessage mCreateRequest;
        private long mPosReserved;
        File localFile;
        private static final AtomicInteger atomic=new AtomicInteger(0);
        private static final String pidname;
        private static final int pidindex;
        private static final String basedir=Configuration.get(PropertyKey.WORKER_TIERED_STORE_LEVEL0_DIRS_PATH);
        static{
            RuntimeMXBean runtime= ManagementFactory.getRuntimeMXBean();
            pidname=runtime.getName();
            pidindex=pidname.indexOf('@');
        }

        AlluxioWritableBlock2(AlluxioFile file,WorkerNetAddress address,long id_,long size_,OutStreamOptions options)throws IOException{
            super(id_,size_);
            StringBuilder builder=new StringBuilder(basedir+"/alluxio.local.block-");
            builder.append(Long.toHexString(System.currentTimeMillis()));
            builder.append('-');
            builder.append(pidname,0,pidindex);
            builder.append('-').append(atomic.getAndIncrement());
            blockPath=builder.toString();
            localFile=new File(blockPath);
            rpcContext=closer.register(new AlluxioRPCContext(file.alluxioFileSystem.fileSystemContext,address));
            ///ClosableResource resource=super.closer.register(new ClosableResource(file.alluxioFileSystem.fileSystemContext,address));
            mCreateRequest = new ProtoMessage(
                    Protocol.LocalBlockCreateRequest.newBuilder().setBlockId(id_)
                            .setTier(options.getWriteTier()).setSpaceToReserve(FILE_BUFFER_BYTES).build());
            ///mNettyRPCContext = NettyRPCContext.defaults().setChannel(resource.channel).setTimeout(WRITE_TIMEOUT_MS);
            mPosReserved=0;
        }

        @Override
        public void cancel()throws IOException {
            if(!isClosed){
                isClosed=true;
                if(localFile.exists()){
                    //noinspection ResultOfMethodCallIgnored
                    localFile.delete();
                }
                closer.close();
            }
        }

        void require(long s) throws IOException {
            if (s <= mPosReserved) {
                return;
            }
            long toReserve = Math.max(s - mPosReserved, FILE_BUFFER_BYTES);
            ///NettyRPC.call(mNettyRPCContext,
            rpcContext.call(new ProtoMessage(mCreateRequest.asLocalBlockCreateRequest().toBuilder().setSpaceToReserve(toReserve)
                            .setOnlyReserveSpace(true).build()),WRITE_TIMEOUT_MS);
            mPosReserved += toReserve;
        }

        @Override
        public int permission() {
            return UnifiedFile.PERMISSION_WRITE;
        }

        @Override
        public void close() throws IOException {
            if(isClosed){
                return;
            }
            isClosed=true;
            try {
                if (!localFile.exists()) {
                    throw new IOException("Unable to find local block");
                }
                super.thisSize = localFile.length();
                ///ProtoMessage message = NettyRPC.call(mNettyRPCContext, mCreateRequest);
                ProtoMessage message=rpcContext.call(mCreateRequest);
                Preconditions.checkState(message.isLocalBlockCreateResponse());
                blockPath = message.asLocalBlockCreateResponse().getPath();
                require(thisSize);
                //noinspection ResultOfMethodCallIgnored
                localFile.renameTo(new File(blockPath));
                Protocol.LocalBlockCompleteRequest request =
                        Protocol.LocalBlockCompleteRequest.newBuilder().setBlockId(thisId).build();
                ///NettyRPC.call(mNettyRPCContext, new ProtoMessage(request));
                rpcContext.call(new ProtoMessage(request),WRITE_TIMEOUT_MS);
            }finally {
                closer.close();
            }
        }
    }

    static class AlluxioWritableBlock extends AlluxioBlock{

        ///private NettyRPCContext mNettyRPCContext;
        private ProtoMessage mCreateRequest;
        private long mPosReserved;
        boolean needCloseOrCancel;

        AlluxioWritableBlock(AlluxioFile file, WorkerNetAddress address,
                             long id_, long size_, OutStreamOptions options) throws IOException {
            super(id_, size_);
            ///ClosableResource resource=super.closer.register(new ClosableResource(file.alluxioFileSystem.fileSystemContext,address));
            mCreateRequest = new ProtoMessage(
                    Protocol.LocalBlockCreateRequest.newBuilder().setBlockId(id_)
                            .setTier(options.getWriteTier()).setSpaceToReserve(FILE_BUFFER_BYTES).build());
            ///mNettyRPCContext = NettyRPCContext.defaults().setChannel(resource.channel).setTimeout(WRITE_TIMEOUT_MS);
            rpcContext=closer.register(new AlluxioRPCContext(file.alluxioFileSystem.fileSystemContext,address));
            ///ProtoMessage message = NettyRPC.call(mNettyRPCContext, mCreateRequest);
            ProtoMessage message=rpcContext.call(mCreateRequest,WRITE_TIMEOUT_MS);
            Preconditions.checkState(message.isLocalBlockCreateResponse());
            super.blockPath = message.asLocalBlockCreateResponse().getPath();
            mPosReserved=0;
            needCloseOrCancel=true;
        }

        AlluxioWritableBlock(AlluxioRPCContext rpc,long id_,long size_,OutStreamOptions options)throws IOException{
            super(id_,size_);
            mCreateRequest = new ProtoMessage(
                    Protocol.LocalBlockCreateRequest.newBuilder().setBlockId(id_)
                            .setTier(options.getWriteTier()).setSpaceToReserve(FILE_BUFFER_BYTES).build());
            rpcContext=rpc;
            ProtoMessage message=rpc.call(mCreateRequest,WRITE_TIMEOUT_MS);
            Preconditions.checkState(message.isLocalBlockCreateResponse());
            blockPath = message.asLocalBlockCreateResponse().getPath();
            mPosReserved=0;
            needCloseOrCancel=true;
        }

        @Override
        public void cancel() throws IOException {
            if (isClosed){
                return;
            }
            isClosed=true;
            try {
                ///NettyRPC.call(mNettyRPCContext,
                rpcContext.call(new ProtoMessage(
                        Protocol.LocalBlockCompleteRequest.newBuilder().setBlockId(super.thisId).setCancel(true)
                                .build()),WRITE_TIMEOUT_MS);
            } catch (Exception e) {
                closer.rethrow(e);
            } finally {
                closer.close();
            }
        }

        void require(long s) throws IOException {
            if (s <= mPosReserved) {
                return;
            }
            long toReserve = Math.max(s - mPosReserved, FILE_BUFFER_BYTES);
            ///NettyRPC.call(mNettyRPCContext,
            rpcContext.call(new ProtoMessage(
                    mCreateRequest.asLocalBlockCreateRequest().toBuilder().setSpaceToReserve(toReserve)
                            .setOnlyReserveSpace(true).build()),WRITE_TIMEOUT_MS);
            mPosReserved += toReserve;
        }

        @Override
        public int permission() {
            return UnifiedFile.PERMISSION_WRITE;
        }

        private void forceClose()throws IOException{
            if(isClosed){
                return;
            }
            File file=new File(super.blockPath);
            super.thisSize=file.length();
            isClosed=true;
            try {
                require(thisSize);
                Protocol.LocalBlockCompleteRequest request =
                        Protocol.LocalBlockCompleteRequest.newBuilder().setBlockId(thisId).build();
                ///NettyRPC.call(mNettyRPCContext, new ProtoMessage(request));
                rpcContext.call(new ProtoMessage(request),WRITE_TIMEOUT_MS);
            }finally {
                closer.close();
            }
        }

        @Override
        public void close() throws IOException {
            if (needCloseOrCancel) {
                forceClose();
            } else {
                cancel();
            }
        }
    }
}
